"""教学用途的 VPN 客户端实现。

该客户端会连接启用 TLS 的 `vpn_server.py` 监听器，并请求其建立指向指定上游
主机/端口的隧道。一旦服务器确认，客户端会将本地绑定的 TCP 套接字（如
`localhost:1080`）中的字节，经由服务器转发到目标主机。

本示例旨在通过简单的 TCP 转发模型演示 VPN 隧道的核心机制，代码保持保守并
配有注释，方便自学。"""

from __future__ import annotations

import argparse
import asyncio
import json
import logging
import ssl
from dataclasses import dataclass
from typing import Optional


@dataclass
class ClientConfig:
    """VPN 客户端的运行时配置。"""

    local_host: str
    local_port: int
    server_host: str
    server_port: int
    auth_token: str
    target_host: str
    target_port: int
    cafile: Optional[str]
    certfile: Optional[str]
    keyfile: Optional[str]


async def open_server_connection(config: ClientConfig) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]:
    """建立到 VPN 服务器的 TLS 连接。"""

    ssl_ctx = ssl.create_default_context(ssl.Purpose.SERVER_AUTH)

    if config.cafile:
        ssl_ctx.load_verify_locations(config.cafile)
    else:
        # 如果未提供 CA 文件则使用系统默认信任链。若是自签名部署需显式指定 `--ca-cert`。
        ssl_ctx.check_hostname = False
        ssl_ctx.verify_mode = ssl.CERT_NONE

    if config.certfile and config.keyfile:
        ssl_ctx.load_cert_chain(certfile=config.certfile, keyfile=config.keyfile)

    reader, writer = await asyncio.open_connection(
        config.server_host,
        config.server_port,
        ssl=ssl_ctx,
    )

    control_payload = json.dumps(
        {
            "token": config.auth_token,
            "host": config.target_host,
            "port": config.target_port,
        }
    ).encode("utf-8") + b"\n"

    writer.write(control_payload)
    await writer.drain()

    response_line = await asyncio.wait_for(reader.readline(), timeout=10)
    response = json.loads(response_line.decode("utf-8"))

    if response.get("status") != "ok":
        raise ConnectionError(f"Server refused tunnel: {response}")

    return reader, writer


async def handle_local_client(
    local_reader: asyncio.StreamReader,
    local_writer: asyncio.StreamWriter,
    config: ClientConfig,
) -> None:
    """处理来自本地应用的 TCP 连接并为其建立隧道。"""

    peername = local_writer.get_extra_info("peername")
    logging.info("Accepted local app connection from %s", peername)

    try:
        server_reader, server_writer = await open_server_connection(config)

        await asyncio.gather(
            relay_stream(local_reader, server_writer),
            relay_stream(server_reader, local_writer),
        )
    except Exception as err:
        logging.warning("Tunnel failed for %s: %s", peername, err)
    finally:
        try:
            local_writer.close()
            await local_writer.wait_closed()
        except Exception:  # pragma: no cover - 尽力完成清理
            logging.debug("Error closing local writer", exc_info=True)


async def relay_stream(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
    """在两个 asyncio 流之间转发字节数据。"""

    try:
        while True:
            chunk = await reader.read(65536)
            if not chunk:
                break
            writer.write(chunk)
            await writer.drain()
    except asyncio.CancelledError:
        raise
    except Exception as err:  # pragma: no cover - 调试辅助
        logging.debug("Relay aborted because of %s", err, exc_info=True)
    finally:
        try:
            writer.close()
            await writer.wait_closed()
        except Exception:
            logging.debug("Stream close raised", exc_info=True)


async def start_local_listener(config: ClientConfig) -> None:
    """监听本地 TCP 客户端，并通过 VPN 服务器转发数据。"""

    server = await asyncio.start_server(
        lambda r, w: handle_local_client(r, w, config),
        host=config.local_host,
        port=config.local_port,
    )

    bind_addresses = ", ".join(str(sock.getsockname()) for sock in server.sockets or [])
    logging.info("Local forwarder listening on %s", bind_addresses)

    async with server:
        await server.serve_forever()


def parse_args() -> ClientConfig:
    parser = argparse.ArgumentParser(description="教学用途的 VPN 客户端")
    parser.add_argument("--local-host", default="127.0.0.1", help="本地应用接入所使用的网络接口")
    parser.add_argument("--local-port", type=int, default=1080, help="本地应用接入的端口")
    parser.add_argument("--server-host", required=True, help="VPN 服务器的主机名或 IP")
    parser.add_argument("--server-port", type=int, default=8443, help="VPN 服务器的端口")
    parser.add_argument("--auth-token", required=True, help="共享的认证令牌")
    parser.add_argument("--target-host", required=True, help="VPN 服务器可访问的目标主机")
    parser.add_argument("--target-port", required=True, type=int, help="目标主机端口")
    parser.add_argument("--ca-cert", dest="cafile", help="用于验证 VPN 服务器证书的 CA 证书包")
    parser.add_argument("--cert", dest="certfile", help="客户端 TLS 证书文件 (PEM)")
    parser.add_argument("--key", dest="keyfile", help="客户端 TLS 私钥文件 (PEM)")
    parser.add_argument("--log-level", default="INFO", help="Python 日志等级（默认：INFO）")

    args = parser.parse_args()

    logging.basicConfig(
        level=getattr(logging, args.log_level.upper(), logging.INFO),
        format="%(asctime)s | %(levelname)s | %(message)s",
    )

    return ClientConfig(
        local_host=args.local_host,
        local_port=args.local_port,
        server_host=args.server_host,
        server_port=args.server_port,
        auth_token=args.auth_token,
        target_host=args.target_host,
        target_port=args.target_port,
        cafile=args.cafile,
        certfile=args.certfile,
        keyfile=args.keyfile,
    )


def main() -> None:
    config = parse_args()

    try:
        asyncio.run(start_local_listener(config))
    except KeyboardInterrupt:
        logging.info("Interrupted by user. Exiting.")


if __name__ == "__main__":
    main()
